{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I9gUzvnVPCoy"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\").\n",
        "\n",
        "# MNIST, with TensorFlow 2.0"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Vm_J8cSKPCo1"
      },
      "source": [
        "## This notebook is still under construction! Please come back later."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "NRXR0hJKPCo2"
      },
      "source": [
        "This notebook trains a simple MNIST model, demonstrating a basic workflow using TensorFlow 2.0 APIs.\n",
        "\n",
        "The basic workflow consists of:\n",
        "\n",
        "- defining a model\n",
        "- preprocessing your data into a tf.data.Dataset\n",
        "- training over a dataset\n",
        "  - using tf.GradientTape to compute gradients\n",
        "  - using stateful tf.keras.metrics.* APIs to collect metrics of interest\n",
        "  - logging those metrics with tf.summary.* APIs so that they can be viewed in TensorBoard\n",
        "  - using tf.train.Checkpoint to save and restore weights\n",
        "- export a SavedModel using tf.saved_model (This SavedModel is a portable representation of the model, and can be imported into C++, JS, Python without knowledge of the original TensorFlow code.)\n",
        "- reimport that SavedModel and demonstrate its usage in Python."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "ZJ5o2Jq6PCo8"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function\n",
        "\n",
        "import os\n",
        "import time\n",
        "\n",
        "import numpy as np\n",
        "!pip install tf-nightly-2.0-preview\n",
        "import tensorflow as tf\n",
        "\n",
        "#TODO(brianklee): remove these when new modules are exported.\n",
        "from tensorflow.python.ops import summary_ops_v2\n",
        "\n",
        "print(tf.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "c3iafiz9PCpA"
      },
      "source": [
        "## Define a convolution-based model, using Keras APIs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Z_cEHfTdPCpB"
      },
      "outputs": [],
      "source": [
        "def create_model():\n",
        "  # Assumes data_format == 'channel_last'.\n",
        "  # See https://www.tensorflow.org/performance/performance_guide#data_formats\n",
        "\n",
        "  l = tf.keras.layers\n",
        "  max_pool = l.MaxPooling2D((2, 2), (2, 2), padding='same')\n",
        "  # The model consists of a sequential chain of layers, so tf.keras.Sequential\n",
        "  # (a subclass of tf.keras.Model) makes for a compact description.\n",
        "  return tf.keras.Sequential([\n",
        "      l.Reshape(\n",
        "          target_shape=[28, 28, 1],\n",
        "          input_shape=(28, 28,)),\n",
        "      l.Conv2D(2, 5, padding='same', activation=tf.nn.relu),\n",
        "      max_pool,\n",
        "      l.Conv2D(4, 5, padding='same', activation=tf.nn.relu),\n",
        "      max_pool,\n",
        "      l.Flatten(),\n",
        "      l.Dense(32, activation=tf.nn.relu),\n",
        "      l.Dropout(0.4),\n",
        "      l.Dense(10)])\n",
        "\n",
        "# Define a loss function and accuracy function\n",
        "def compute_loss(logits, labels):\n",
        "  return tf.reduce_mean(\n",
        "      tf.nn.sparse_softmax_cross_entropy_with_logits(\n",
        "          logits=logits, labels=labels))\n",
        "\n",
        "\n",
        "def compute_accuracy(logits, labels):\n",
        "  predictions = tf.argmax(logits, axis=1, output_type=tf.int64)\n",
        "  labels = tf.cast(labels, tf.int64)\n",
        "  return tf.reduce_mean(\n",
        "      tf.cast(tf.equal(predictions, labels), dtype=tf.float32))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YsaXPR6OPCpE"
      },
      "outputs": [],
      "source": [
        "# Create the model and optimizer\n",
        "model = create_model()\n",
        "\n",
        "optimizer = tf.keras.optimizers.SGD(learning_rate=0.01, momentum=0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7bgV8B1wPCpJ"
      },
      "source": [
        "## Download and create tf.data.Datasets\n",
        "\n",
        "Let's load the MNIST dataset into TF datasets. This lets us use all sorts of useful transformations like batching and shuffling. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "zxmeEGyhPCpL"
      },
      "outputs": [],
      "source": [
        "# Set up datasets\n",
        "def mnist_datasets():\n",
        "  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "  # Numpy defaults to dtype=float64; TF defaults to float32. Stick with float32.\n",
        "  x_train, x_test = x_train / np.float32(255), x_test / np.float32(255)\n",
        "  y_train, y_test = y_train.astype(np.int64), y_test.astype(np.int64)\n",
        "  train_dataset = tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "  test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test))\n",
        "  return train_dataset, test_dataset\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YtB65OO0PCpP"
      },
      "outputs": [],
      "source": [
        "train_ds, test_ds = mnist_datasets()\n",
        "train_ds = train_ds.shuffle(60000).batch(100)\n",
        "test_ds = test_ds.batch(100)\n",
        "\n",
        "print('Dataset will yield tensors of the following shape: {}'.format(train_ds.output_shapes))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uB_C1r9PPCpU"
      },
      "source": [
        "## Configure training\n",
        "\n",
        "Our train() function does a few things. It iterates over our training dataset, computing the gradients for each batch and then applying them to the model variables. Along the way, we periodically output summaries.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "G8EfprJ1PCpU"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(model, optimizer, images, labels):\n",
        "  # Record the operations used to compute the loss, so that the gradient\n",
        "  # of the loss with respect to the variables can be computed.\n",
        "  with tf.GradientTape() as tape:\n",
        "    logits = model(images, training=True)\n",
        "    loss = compute_loss(logits, labels)\n",
        "    accuracy = compute_accuracy(logits, labels)\n",
        "  grads = tape.gradient(loss, model.variables)\n",
        "  optimizer.apply_gradients(zip(grads, model.variables))\n",
        "  return loss, accuracy\n",
        "\n",
        "\n",
        "def train(model, optimizer, dataset, log_freq=10):\n",
        "  \"\"\"Trains model on `dataset` using `optimizer`.\"\"\"\n",
        "  start = time.time()\n",
        "  # Metrics are stateful. They accumulate values and return a cumulative\n",
        "  # result when you call .result(). Clear accumulated values with .reset_states()\n",
        "  avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)\n",
        "  avg_accuracy = tf.keras.metrics.Mean('accuracy', dtype=tf.float32)\n",
        "  # Datasets can be iterated over like any other Python iterable.\n",
        "  for images, labels in dataset:\n",
        "    loss, accuracy = train_step(model, optimizer, images, labels)\n",
        "    avg_loss(loss)\n",
        "    avg_accuracy(accuracy)\n",
        "    if tf.equal(optimizer.iterations % log_freq, 0):\n",
        "      summary_ops_v2.scalar('loss', avg_loss.result(), step=optimizer.iterations)\n",
        "      summary_ops_v2.scalar('accuracy', avg_accuracy.result(), step=optimizer.iterations)\n",
        "      avg_loss.reset_states()\n",
        "      avg_accuracy.reset_states()\n",
        "      rate = log_freq / (time.time() - start)\n",
        "      print('Step #%d\\tLoss: %.6f (%d steps/sec)' % (optimizer.iterations, loss, rate))\n",
        "      start = time.time()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6zrAPkdEPCpa"
      },
      "outputs": [],
      "source": [
        "def test(model, dataset, step_num):\n",
        "  \"\"\"Perform an evaluation of `model` on the examples from `dataset`.\"\"\"\n",
        "  avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)\n",
        "  avg_accuracy = tf.keras.metrics.Mean('accuracy', dtype=tf.float32)\n",
        "\n",
        "  for (images, labels) in dataset:\n",
        "    logits = model(images, training=False)\n",
        "    avg_loss(compute_loss(logits, labels))\n",
        "    avg_accuracy(compute_accuracy(logits, labels))\n",
        "  print('Model test set loss: {:0.4f} accuracy: {:0.2f}%'.format(\n",
        "      avg_loss.result(), avg_accuracy.result() * 100))\n",
        "  summary_ops_v2.scalar('loss', avg_loss.result(), step=step_num)\n",
        "  summary_ops_v2.scalar('accuracy', avg_accuracy.result(), step=step_num)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kox8FEeNPCpd"
      },
      "source": [
        "## Configure model directory\n",
        "\n",
        "We'll use one directory to save all of our relevant artifacts (summary logs, checkpoints, SavedModel exports, etc.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "bLwFfkYhPCpe"
      },
      "outputs": [],
      "source": [
        "# Where to save checkpoints, tensorboard summaries, etc.\n",
        "MODEL_DIR = '/tmp/tensorflow/mnist'\n",
        "\n",
        "def apply_clean():\n",
        "  if tf.io.gfile.exists(MODEL_DIR):\n",
        "    print('Removing existing model dir: {}'.format(MODEL_DIR))\n",
        "    tf.io.gfile.rmtree(MODEL_DIR)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "5DUL7OVYPCph"
      },
      "outputs": [],
      "source": [
        "# Optional: wipe the existing directory\n",
        "\n",
        "apply_clean()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "SL51Mdg9PCpj"
      },
      "source": [
        "Let's configure where we'll output our summaries from training. \n",
        "\n",
        "Remember how we called tf.summary.scalar(...) in our train() function? By using our summary_writer in a `with` block, we can catch those generated summaries and direct them to a file.\n",
        "\n",
        "You can see the summaries with `tensorboard --logdir=\u003cmodel_dir\u003e`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YZgxx95-PCpk"
      },
      "outputs": [],
      "source": [
        "train_dir = os.path.join(MODEL_DIR, 'summaries', 'train')\n",
        "test_dir = os.path.join(MODEL_DIR, 'summaries', 'eval')\n",
        "train_summary_writer = summary_ops_v2.create_file_writer(\n",
        "  train_dir, flush_millis=10000)\n",
        "test_summary_writer = summary_ops_v2.create_file_writer(\n",
        "  test_dir, flush_millis=10000, name='test')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pLSsDST8PCpn"
      },
      "source": [
        "## Configure checkpoints\n",
        "\n",
        "The Checkpoint object helps manage which tf.Variables will be saved to and restored from checkpoint files.\n",
        "\n",
        "A checkpoint differs from a SavedModel because it must additionally keep track of training-related state, like momentum variables for a momentum-based optimizer or things like the global step. Additionally, it only stores weights, so you'll need the original code to define the computation using those weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "Xab3feHXPCpp"
      },
      "outputs": [],
      "source": [
        "checkpoint_dir = os.path.join(MODEL_DIR, 'checkpoints')\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt')\n",
        "checkpoint = tf.train.Checkpoint(\n",
        "  model=model, optimizer=optimizer)\n",
        "# Restore variables on creation if a checkpoint exists.\n",
        "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "U0ZNI7rFPCps"
      },
      "source": [
        "## Train\n",
        "\n",
        "Now that we've set up train() and test(), let's create a model and train it for some number of epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "CagzdcoRFrN9"
      },
      "outputs": [],
      "source": [
        "NUM_TRAIN_EPOCHS = 1\n",
        "\n",
        "for i in range(NUM_TRAIN_EPOCHS):\n",
        "  start = time.time()\n",
        "  with train_summary_writer.as_default():\n",
        "    train(model, optimizer, train_ds)\n",
        "  end = time.time()\n",
        "  print('\\nTrain time for epoch #{} ({} total steps): {}'.format(\n",
        "      i + 1, optimizer.iterations, end - start))\n",
        "  with test_summary_writer.as_default():\n",
        "    test(model, test_ds, optimizer.iterations)\n",
        "  checkpoint.save(checkpoint_prefix)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "dc1hV_21PCpw"
      },
      "source": [
        "## Export a SavedModel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "NGCnOqMTPCpy"
      },
      "outputs": [],
      "source": [
        "export_path = os.path.join(MODEL_DIR, 'export')\n",
        "tf.saved_model.save(model, export_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uECIhO7aPCp1"
      },
      "source": [
        "## Restore and run the SavedModel\n",
        "\n",
        "You can restore any SavedModel and call it without reference to the original source code. APIs for importing and transforming SavedModels exist a variety of languages. See the [guide](https://www.tensorflow.org/guide/saved_model) for more."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "go-t0CvMzrSi"
      },
      "outputs": [],
      "source": [
        "def import_and_eval():\n",
        "  restored_model = tf.saved_model.restore(export_path)\n",
        "  _, (x_test, y_test) = tf.keras.datasets.mnist.load_data()\n",
        "  x_test = x_test / np.float32(255)\n",
        "  y_predict = restored_model(x_test)\n",
        "  accuracy = compute_accuracy(y_predict, y_test)\n",
        "  print('Model accuracy: {:0.2f}%'.format(accuracy.numpy() * 100))\n",
        "\n",
        "# TODO(brianklee): Activate after v2 import is implemented.\n",
        "# import_and_eval()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "eNY8eU40PCp8"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "last_runtime": {
        "build_target": "",
        "kind": "local"
      },
      "name": "Untitled2.ipynb",
      "private_outputs": true,
      "provenance": [],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
